msg_tool\scripts\emote/
dref.rs1use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::img::*;
8use anyhow::Result;
9use emote_psb::PsbReader;
10use std::collections::HashMap;
11use std::io::Read;
12use std::path::{Path, PathBuf};
13use url::Url;
14
15#[derive(Debug)]
16pub struct DrefBuilder {}
18
19impl DrefBuilder {
20 pub fn new() -> Self {
22 Self {}
23 }
24}
25
26impl ScriptBuilder for DrefBuilder {
27 fn default_encoding(&self) -> Encoding {
28 Encoding::Cp932
29 }
30
31 fn build_script(
32 &self,
33 buf: Vec<u8>,
34 filename: &str,
35 encoding: Encoding,
36 _archive_encoding: Encoding,
37 config: &ExtraConfig,
38 archive: Option<&Box<dyn Script>>,
39 ) -> Result<Box<dyn Script>> {
40 Ok(Box::new(Dref::new(
41 buf, encoding, filename, config, archive,
42 )?))
43 }
44
45 fn extensions(&self) -> &'static [&'static str] {
46 &["dref"]
47 }
48
49 fn script_type(&self) -> &'static ScriptType {
50 &ScriptType::EmoteDref
51 }
52
53 fn is_image(&self) -> bool {
54 true
55 }
56}
57
58struct Dpak {
59 psb: VirtualPsbFixed,
60}
61
62struct OffsetData {
63 left: u32,
64 top: u32,
65}
66
67impl Dpak {
68 pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
69 let f = std::fs::File::open(path)?;
70 let mut f = std::io::BufReader::new(f);
71 let mut psb = PsbReader::open_psb(&mut f)
72 .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK: {:?}", e))?;
73 let psb = psb
74 .load()
75 .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK: {:?}", e))?;
76 let psb = psb.to_psb_fixed();
77 Ok(Self { psb })
78 }
79
80 pub fn load_from_data(data: &[u8]) -> Result<Self> {
81 let mut psb = PsbReader::open_psb(MemReaderRef::new(data))
82 .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK data: {:?}", e))?;
83 let psb = psb
84 .load()
85 .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK data: {:?}", e))?;
86 let psb = psb.to_psb_fixed();
87 Ok(Self { psb })
88 }
89
90 pub fn load_image(&self, name: &str) -> Result<(ImageData, Option<OffsetData>)> {
91 let root = self.psb.root();
92 let rid = root[name]
93 .resource_id()
94 .ok_or_else(|| anyhow::anyhow!("Resource ID for image '{}' not found in DPAK", name))?
95 as usize;
96 if rid >= self.psb.resources().len() {
97 return Err(anyhow::anyhow!(
98 "Resource ID {} out of bounds for DPAK with {} resources",
99 rid,
100 self.psb.resources().len()
101 ));
102 }
103 let resource = &self.psb.resources()[rid];
104 Self::load_png(&resource)
105 }
106
107 fn load_png(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
108 let mut img = load_png(MemReaderRef::new(&data))?;
109 match img.color_type {
110 ImageColorType::Rgb => {
111 convert_rgb_to_rgba(&mut img)?;
112 }
113 _ => {}
114 }
115 Ok((
116 img,
117 Self::try_read_offset_from_png(MemReaderRef::new(&data))?,
118 ))
119 }
120
121 fn try_read_offset_from_png(mut data: MemReaderRef) -> Result<Option<OffsetData>> {
122 data.pos = 8; data.pos += 8; data.pos += 17; loop {
126 let chunk_size = data.read_u32_be()?;
127 let mut chunk_type = [0u8; 4];
128 data.read_exact(&mut chunk_type)?;
129 if &chunk_type == b"IDAT" || &chunk_type == b"IEND" {
130 break;
131 }
132 if &chunk_type == b"oFFs" {
133 let x = data.read_u32_be()?;
134 let y = data.read_u32_be()?;
135 if data.read_u8()? == 0 {
136 return Ok(Some(OffsetData { left: x, top: y }));
137 }
138 }
139 data.pos += chunk_size as usize + 4; }
141 Ok(None)
142 }
143}
144
145#[derive(Default)]
146struct DpakLoader {
147 map: HashMap<String, Dpak>,
148}
149
150impl DpakLoader {
151 pub fn load_image(
152 &mut self,
153 dir: &Path,
154 dpak: &str,
155 filename: &str,
156 ) -> Result<(ImageData, Option<OffsetData>)> {
157 let dpak = match self.map.get(dpak) {
158 Some(d) => d,
159 None => {
160 let path = dir.join(dpak);
161 let ndpak = Dpak::new(&path)?;
162 self.map.insert(dpak.to_string(), ndpak);
163 self.map.get(dpak).unwrap()
164 }
165 };
166 dpak.load_image(filename)
167 }
168
169 pub fn load_archives(&mut self, in_archives: &HashMap<String, Vec<u8>>) -> Result<()> {
170 for (name, data) in in_archives.iter() {
171 if !self.map.contains_key(name) {
172 let dpak = Dpak::load_from_data(data)?;
173 self.map.insert(name.clone(), dpak);
174 }
175 }
176 Ok(())
177 }
178}
179
180pub struct Dref {
182 urls: Vec<Url>,
183 dir: PathBuf,
184 in_archives: HashMap<String, Vec<u8>>,
185}
186
187impl std::fmt::Debug for Dref {
188 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189 f.debug_struct("Dref")
190 .field("urls", &self.urls)
191 .field("dir", &self.dir)
192 .finish()
193 }
194}
195
196impl Dref {
197 pub fn new(
205 buf: Vec<u8>,
206 encoding: Encoding,
207 filename: &str,
208 _config: &ExtraConfig,
209 archive: Option<&Box<dyn Script>>,
210 ) -> Result<Self> {
211 let text = decode_with_bom_detect(encoding, &buf, true)?.0;
212 let mut urls = Vec::new();
213 for text in text.lines() {
214 let text = text.trim();
215 if text.is_empty() {
216 continue;
217 }
218 urls.push(Url::parse(text)?);
219 }
220 let path = Path::new(filename);
221 let dir = if let Some(parent) = path.parent() {
222 parent.to_path_buf()
223 } else {
224 PathBuf::from(".")
225 };
226 if urls.is_empty() {
227 return Err(anyhow::anyhow!("No URLs found in DREF file: {}", filename));
228 }
229 for u in urls.iter() {
230 if u.scheme() != "psb" {
231 return Err(anyhow::anyhow!(
232 "Invalid URL scheme in DREF file: {} (expected 'psb')",
233 u
234 ));
235 }
236 }
237 let mut in_archives = HashMap::new();
238 if let Some(archive) = archive {
239 if archive.is_archive() {
240 for url in urls.iter() {
241 let filename = url.domain().ok_or(anyhow::anyhow!(
242 "Invalid URL in DREF file: {} (missing domain)",
243 url
244 ))?;
245 if let Ok(mut content) = archive.open_file_by_name(filename, true) {
246 in_archives.insert(filename.to_string(), content.data()?);
247 }
248 }
249 }
250 }
251 Ok(Self {
252 urls,
253 dir,
254 in_archives,
255 })
256 }
257}
258
259impl Script for Dref {
260 fn default_output_script_type(&self) -> OutputScriptType {
261 OutputScriptType::Json
262 }
263
264 fn default_format_type(&self) -> FormatOptions {
265 FormatOptions::None
266 }
267
268 fn is_image(&self) -> bool {
269 true
270 }
271
272 fn export_image(&self) -> Result<ImageData> {
273 let mut loader = DpakLoader::default();
274 loader.load_archives(&self.in_archives)?;
275 let base_url = &self.urls[0];
276 let dpak = base_url.domain().ok_or(anyhow::anyhow!(
277 "Invalid URL in DREF file: {} (missing domain)",
278 base_url
279 ))?;
280 let (mut base_img, base_offset) =
281 loader.load_image(&self.dir, dpak, base_url.path().trim_start_matches("/"))?;
282 if let Some(o) = base_offset {
283 eprintln!("WARN: Base image offset: left={}, top={}", o.left, o.top);
284 crate::COUNTER.inc_warning();
285 }
286 for url in &self.urls[1..] {
287 let dpak = url.domain().ok_or(anyhow::anyhow!(
288 "Invalid URL in DREF file: {} (missing domain)",
289 url
290 ))?;
291 let (img, img_offset) =
292 loader.load_image(&self.dir, dpak, url.path().trim_start_matches("/"))?;
293 let (top, left) = match img_offset {
294 Some(o) => (o.top, o.left),
295 None => (0, 0),
296 };
297 draw_on_img_with_opacity(&mut base_img, &img, left, top, 0xff)?;
298 }
299 Ok(base_img)
300 }
301}